import streamlit as st
import os
import pandas as pd
from utils.helpers import get_products_from_data_dir, get_product_stats

def show_home():
    """显示首页"""
    
    st.title("DefectRAG - 产品表面缺陷检测系统")
    
    st.markdown("""
    ## 系统简介
    
    DefectRAG是一个基于CN_CLIP + RAG + VLLM的产品表面缺陷检测系统。本系统利用大型视觉语言模型和检索增强生成技术，实现高精度的产品表面缺陷检测和分析。
    
    ### 主要功能
    
    1. **知识库管理**：构建和管理产品缺陷知识库
    2. **缺陷检测**：检测图片中的产品表面缺陷并输出结构化结果
    3. **对话模式**：与系统进行自然语言交互，询问产品缺陷相关问题
    4. **系统测评**：对系统性能进行全面评估，生成详细报告
    
    ### 技术架构
    
    - **特征提取**：CN_CLIP (Chinese Contrastive Language-Image Pre-training)
    - **知识检索**：基于FAISS的高效向量检索
    - **推理生成**：基于Ollama的VLLM (Qwen2.5-VL) 推理
    """)
    
    st.markdown("---")
    
    # 显示数据集信息
    st.header("数据集概览")
    
    data_dir = "./data"
    products = get_products_from_data_dir(data_dir)
    
    if not products:
        st.warning("未找到可用的产品数据集。请检查数据目录。")
        return
    
    # 收集所有产品的统计信息
    all_stats = []
    for product in products:
        stats = get_product_stats(data_dir, product)
        all_stats.append({
            "产品": product,
            "训练图片数": stats["total_train_images"],
            "测试图片数": stats["total_test_images"],
            "缺陷类型数": len(stats["train"]) - (1 if "good" in stats["train"] else 0)
        })
    
    # 显示统计表格
    df = pd.DataFrame(all_stats)
    st.dataframe(df, use_container_width=True)
    
    # 选择查看详细信息
    selected_product = st.selectbox("选择产品查看详细信息", [""] + products)
    
    if selected_product:
        stats = get_product_stats(data_dir, selected_product)
        
        col1, col2 = st.columns(2)
        
        # 训练集信息
        with col1:
            st.subheader("训练集")
            train_data = []
            for defect_type, count in stats["train"].items():
                train_data.append({"类别": defect_type, "图片数量": count})
            
            train_df = pd.DataFrame(train_data)
            st.dataframe(train_df, use_container_width=True)
        
        # 测试集信息
        with col2:
            st.subheader("测试集")
            test_data = []
            for defect_type, count in stats["test"].items():
                test_data.append({"类别": defect_type, "图片数量": count})
            
            test_df = pd.DataFrame(test_data)
            st.dataframe(test_df, use_container_width=True)
    
    st.markdown("---")
    
    st.header("使用指南")
    
    st.markdown("""
    1. **知识库管理**：首先创建产品的缺陷知识库
       - 选择CN_CLIP模型
       - 从本地文件夹或上传图片构建知识库
       
    2. **缺陷检测**：上传图片进行缺陷检测
       - 选择产品和检测模型
       - 系统会返回结构化的检测结果
       
    3. **对话模式**：与系统进行自然语言交互
       - 上传图片并提出问题
       - 系统会结合图像内容回答问题
       
    4. **系统测评**：评估系统性能
       - 选择测试的产品
       - 运行测评并生成报告
    """)
    
    st.markdown("---")
    
    st.info("开发者：DefectRAG团队") 